package com.longyou.gateway.filter;

import java.util.List;
import java.util.regex.Pattern;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

@Slf4j
public class ShellInjectionWebFilter implements WebFilter {

    // 更精确的正则表达式，只拦截真正可疑的命令组合
    private static final Pattern SHELL_PATTERN = Pattern.compile(
        "(;\\s*(sh|bash|zsh|ksh|csh|tcsh|dash|chmod|wget|curl|nc|telnet|perl|python|ruby)\\b|" +
            "&&\\s*(sh|bash|zsh|ksh|csh|tcsh|dash|chmod|wget|curl|nc|telnet|perl|python|ruby)\\b|" +
            "\\|\\|\\s*(sh|bash|zsh|ksh|csh|tcsh|dash|chmod|wget|curl|nc|telnet|perl|python|ruby)\\b|" +
            "\\$\\(.*\\)|`.*`|" +
            "\\b(sh|bash|zsh|ksh|csh|tcsh|dash)\\s+.*\\b(chmod|wget|curl|nc|telnet|perl|python|ruby)\\b)",
        Pattern.CASE_INSENSITIVE
    );

    // 白名单：常见的正常请求头值
    private static final Pattern WHITELIST_PATTERN = Pattern.compile(
        "(Mozilla|Chrome|Safari|Firefox|Edge|Windows|Mac|Linux|Android|iOS|" +
            "application/json|text/plain|gzip|br|zh-CN|en-US|" +
            "cloudflare|same-origin|cors|https|http|" +
            "\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}|" + // IP地址
            "[0-9a-fA-F:]+)" // IPv6地址
    );

    @Override
    public @NotNull Mono<Void> filter(ServerWebExchange exchange, @NotNull WebFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();

        // 检查查询参数
        if (hasMaliciousQueryParams(request)) {
            log.warn("Malicious query parameter detected: {}", request.getURI());
            exchange.getResponse().setStatusCode(HttpStatus.BAD_REQUEST);
            return exchange.getResponse().setComplete();
        }

        // 检查请求头（更精确的检测）
        if (hasMaliciousHeaders(request)) {
            log.warn("Malicious header detected: {}", request.getURI());
            exchange.getResponse().setStatusCode(HttpStatus.BAD_REQUEST);
            return exchange.getResponse().setComplete();
        }

        // 如果没有问题，继续执行链
        return chain.filter(exchange);
    }

    private boolean hasMaliciousQueryParams(ServerHttpRequest request) {
        return request.getQueryParams().values().stream()
            .flatMap(List::stream)
            .filter(value -> value != null && value.length() > 5) // 只检测长度大于5的参数
            .anyMatch(this::isMaliciousValue);
    }

    private boolean hasMaliciousHeaders(ServerHttpRequest request) {
        return request.getHeaders().entrySet().stream()
            .filter(entry -> {
                String headerName = entry.getKey().toLowerCase();
                // 只检查可能包含用户输入的请求头
                return headerName.contains("user") ||
                    headerName.contains("agent") ||
                    headerName.contains("referer") ||
                    headerName.contains("origin") ||
                    headerName.contains("x-");
            })
            .flatMap(entry -> entry.getValue().stream())
            .filter(value -> value != null && value.length() > 10) // 只检测长度大于10的请求头值
            .anyMatch(this::isMaliciousValue);
    }

    private boolean isMaliciousValue(String value) {
        if (value == null || value.trim().isEmpty()) {
            return false;
        }

        // 如果值在白名单中，直接放行
        if (WHITELIST_PATTERN.matcher(value).find()) {
            return false;
        }

        // 检查是否包含可疑的shell命令模式
        return SHELL_PATTERN.matcher(value).find();
    }
}